Skip to content

Refactor WAN loading logic: granular flags & boilerplate cleanup#419

Open
ninatu wants to merge 1 commit into
mainfrom
ninatu/wan_load_params
Open

Refactor WAN loading logic: granular flags & boilerplate cleanup#419
ninatu wants to merge 1 commit into
mainfrom
ninatu/wan_load_params

Conversation

@ninatu

@ninatu ninatu commented Jun 16, 2026

Copy link
Copy Markdown
Collaborator

Replace 'vae_only' and 'transformer_only' flags in WAN pipelines and checkpointers with granular 'load_*' boolean flags. Centralize repetitive loading and instantiation logic (from_pretrained, from_checkpoint, etc.) into the base classes.

@ninatu ninatu requested a review from entrpn as a code owner June 16, 2026 09:05
@github-actions

Copy link
Copy Markdown

@ninatu ninatu force-pushed the ninatu/wan_load_params branch 3 times, most recently from 1e21945 to 135542b Compare June 16, 2026 13:42
@github-actions

Copy link
Copy Markdown

🤖 Hi @csgoogle, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions

Copy link
Copy Markdown

🤖 I'm sorry @csgoogle, but I was unable to process your request. Please see the logs for more details.

@ninatu

ninatu commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator Author

Hi @entrpn! Not sure why gemini review is failing, is there anything I can do about it?

@github-actions

Copy link
Copy Markdown

🤖 Hi @entrpn, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📋 Review Summary

This pull request presents an exceptionally high-quality refactoring of the WAN loading logic. By introducing granular boolean loading flags and centralizing the boilerplate instantiation, loading, and checkpoint management from 5 separate subclasses into base classes, it dramatically simplifies codebase complexity and eliminates duplicate code.

🔍 General Feedback

  • Boilerplate Reduction: Consolidating load_checkpoint and load_diffusers_checkpoint logic into WanCheckpointer reduces significant duplicate logic across 5 checkpointer files, which is a major win for maintainability.
  • Granular Loading Control: The replacement of vae_only and transformer_only with granular flags (load_vae, load_text_encoder, load_transformer, load_scheduler) provides precise control over pipeline resources, allowing users to optimize resource initialization dynamically.
  • Exhaustive Unit Tests: The addition of scenarios testing all permutations of the load flags in wan_vace_pipeline_test.py is exemplary and ensures robust regression coverage.

Comment thread src/maxdiffusion/pipelines/wan/wan_pipeline.py Outdated
@github-actions

Copy link
Copy Markdown

🤖 I'm sorry @entrpn, but I was unable to process your request. Please see the logs for more details.

Replace  'vae_only' and 'transformer_only' flags in WAN pipelines and
checkpointers with granular 'load_*' boolean flags. Centralize repetitive
loading and instantiation logic (from_pretrained, from_checkpoint, etc.)
into the base classes.
@ninatu ninatu force-pushed the ninatu/wan_load_params branch from 135542b to ef5af09 Compare June 17, 2026 17:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants